#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <geometry_msgs/PoseStamped.h>
#include <pcl/point_types.h>
#include <pcl/registration/ndt.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/io/pcd_io.h>
#include <tf2_ros/transform_broadcaster.h>
#include <tf2/LinearMath/Transform.h>
#include <tf2_eigen/tf2_eigen.h>

class NDTLocalizationTester {
private:
    ros::NodeHandle nh_;
    ros::NodeHandle private_nh_;
    ros::Subscriber points_sub_;
    ros::Publisher pose_pub_;
    ros::Publisher map_pub_;
    tf2_ros::TransformBroadcaster tf_broadcaster_;

    // NDT parameters
    double ndt_resolution_;
    double transformation_epsilon_;
    double step_size_;
    int max_iterations_;
    double voxel_leaf_size_;
    std::string map_file_;

    // PCL NDT object
    pcl::NormalDistributionsTransform<pcl::PointXYZ, pcl::PointXYZ> ndt_;
    pcl::PointCloud<pcl::PointXYZ>::Ptr map_cloud_;
    Eigen::Matrix4f initial_pose_;
    bool has_initial_pose_;

public:
    NDTLocalizationTester() : private_nh_("~"), map_cloud_(new pcl::PointCloud<pcl::PointXYZ>), has_initial_pose_(false) {
        // Load parameters
        private_nh_.param("ndt_resolution", ndt_resolution_, 1.0);
        private_nh_.param("transformation_epsilon", transformation_epsilon_, 0.01);
        private_nh_.param("step_size", step_size_, 0.1);
        private_nh_.param("max_iterations", max_iterations_, 30);
        private_nh_.param("voxel_leaf_size", voxel_leaf_size_, 0.5);
        private_nh_.param<std::string>("map_file", map_file_, "");

        // Initialize NDT
        ndt_.setTransformationEpsilon(transformation_epsilon_);
        ndt_.setStepSize(step_size_);
        ndt_.setResolution(ndt_resolution_);
        ndt_.setMaximumIterations(max_iterations_);

        // Load map
        if (pcl::io::loadPCDFile<pcl::PointXYZ>(map_file_, *map_cloud_) == -1) {
            ROS_ERROR("Failed to load map file: %s", map_file_.c_str());
            return;
        }
        ROS_INFO("Loaded map with %lu points", map_cloud_->size());
        ndt_.setInputTarget(map_cloud_);

        // Set initial pose to identity
        initial_pose_ = Eigen::Matrix4f::Identity();

        // Setup ROS communication
        points_sub_ = nh_.subscribe("points_raw", 10, &NDTLocalizationTester::pointsCallback, this);
        pose_pub_ = nh_.advertise<geometry_msgs::PoseStamped>("ndt/current_pose", 10);
        map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("map_cloud", 1, true);

        // Publish map as PointCloud2
        sensor_msgs::PointCloud2 map_msg;
        pcl::toROSMsg(*map_cloud_, map_msg);
        map_msg.header.frame_id = "map";
        map_pub_.publish(map_msg);
    }

    void pointsCallback(const sensor_msgs::PointCloud2::ConstPtr& input_cloud_msg) {
        pcl::PointCloud<pcl::PointXYZ>::Ptr input_cloud(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::fromROSMsg(*input_cloud_msg, *input_cloud);

        // Downsample input cloud
        pcl::VoxelGrid<pcl::PointXYZ> voxel_grid;
        pcl::PointCloud<pcl::PointXYZ>::Ptr filtered_cloud(new pcl::PointCloud<pcl::PointXYZ>);
        voxel_grid.setLeafSize(voxel_leaf_size_, voxel_leaf_size_, voxel_leaf_size_);
        voxel_grid.setInputCloud(input_cloud);
        voxel_grid.filter(*filtered_cloud);

        // Set input cloud for NDT
        ndt_.setInputSource(filtered_cloud);

        // Align
        pcl::PointCloud<pcl::PointXYZ>::Ptr output_cloud(new pcl::PointCloud<pcl::PointXYZ>);
        ndt_.align(*output_cloud, initial_pose_);

        if (ndt_.hasConverged()) {
            initial_pose_ = ndt_.getFinalTransformation();
            double score = ndt_.getFitnessScore();
            ROS_INFO("NDT has converged. Score: %f", score);

            // Publish pose
            geometry_msgs::PoseStamped pose_msg;
            pose_msg.header.frame_id = "map";
            pose_msg.header.stamp = input_cloud_msg->header.stamp;

            Eigen::Affine3d pose_eigen;
            pose_eigen.matrix() = initial_pose_.cast<double>();
            
            geometry_msgs::Transform transform = tf2::eigenToTransform(pose_eigen).transform;
            pose_msg.pose.position.x = transform.translation.x;
            pose_msg.pose.position.y = transform.translation.y;
            pose_msg.pose.position.z = transform.translation.z;
            pose_msg.pose.orientation = transform.rotation;
            pose_pub_.publish(pose_msg);

            // Broadcast transform
            geometry_msgs::TransformStamped transform_stamped;
            transform_stamped.header = pose_msg.header;
            transform_stamped.child_frame_id = "base_link";
            transform_stamped.transform = transform;
            tf_broadcaster_.sendTransform(transform_stamped);
        } else {
            ROS_WARN("NDT did not converge!");
        }
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "test_ndt_localization");
    NDTLocalizationTester tester;
    ros::spin();
    return 0;
} 